﻿using Ndns.ServerHost.Abstractions;
using System;
using System.Collections.Generic;
using Ndns.ServerHost.Abstractions.Services;
using System.Threading.Tasks;
using System.Net;
using System.Threading;
using Ndns.Net;
using Ndns.Packet.Parser;
using Ndns.Packet.Abstractions;

namespace Ndns.Forwarder
{
    class DnsForwarderMiddleware : IMiddleware
    {
        private readonly DnsPacketParser _packetParser = new DnsPacketParser(SharedObjects.ArrayPool);

        private readonly List<IPEndPoint> _udpServers = new List<IPEndPoint>();
        private readonly List<IPEndPoint> _tcpServers = new List<IPEndPoint>();

        private IServiceCollection _services;
        private RequestDelegate _next;
        private ILogger _logger;

        private readonly Random _rand = new Random(Guid.NewGuid().GetHashCode());

        private readonly CancellationTokenSource _disposedToken = new CancellationTokenSource();

        public DnsForwarderMiddleware(IEnumerable<DnsForwarderInfo> forwarders)
        {
            foreach (var forwarder in forwarders)
            {
                switch (forwarder.Protocol)
                {
                    case DnsProtocol.Udp:
                        _udpServers.Add(forwarder.EndPoint);
                        break;
                    case DnsProtocol.Tcp:
                        _tcpServers.Add(forwarder.EndPoint);
                        break;
                    default:
                        // Something is wrong!
                        break;
                }
            }
            if ((_udpServers.Count + _tcpServers.Count) == 0)
            {
                throw new ArgumentException("No valid forwarders provided.", nameof(forwarders));
            }
        }

        public void RegisterServices(IServiceCollection services, RequestDelegate next)
        {
            _services = services;
            _next = next;
            _logger = services.Get<ILoggerProvider>().GetLogger<DnsForwarderMiddleware>();
        }

        private DnsForwarderInfo GetForworder(DnsProtocol protocol)
        {
            if (protocol == DnsProtocol.Tcp)
            {
                if (_tcpServers.Count > 0)
                {
                    return new DnsForwarderInfo(DnsProtocol.Tcp, _tcpServers[_rand.Next(0, _tcpServers.Count)]);
                }
            }
            if (_udpServers.Count > 0)
            {
                return new DnsForwarderInfo(DnsProtocol.Udp, _udpServers[_rand.Next(0, _udpServers.Count)]);
            }
            if (_tcpServers.Count > 0)
            {
                return new DnsForwarderInfo(DnsProtocol.Tcp, _tcpServers[_rand.Next(0, _tcpServers.Count)]);
            }

            return new DnsForwarderInfo(DnsProtocol.Unknown, null);
        }

        public Task InvokeAsync(DnsContext context)
        {
            if (!context.Request.RecursionDesired)
            {
                _logger.LogInfo(context.TraceIdentifier, "Skipping forwarder because RecursionDesired flag it not set.");
                return _next(context);
            }

            DnsForwarderInfo forwarderInfo = GetForworder(context.Connection.Protocol);

            var linkedCts = CancellationTokenSource.CreateLinkedTokenSource(context.RequestAborted, _disposedToken.Token);

            if (forwarderInfo.Protocol == DnsProtocol.Udp)
            {
                _logger.LogInfo(context.TraceIdentifier, $"Forwarding request to [{forwarderInfo.Protocol}]{forwarderInfo.EndPoint}");
                return UdpForwardAsync(context, forwarderInfo.EndPoint, linkedCts.Token);
            }
            else if (forwarderInfo.Protocol == DnsProtocol.Tcp)
            {
                _logger.LogInfo(context.TraceIdentifier, $"Forwarding request to [{forwarderInfo.Protocol}]{forwarderInfo.EndPoint}");
                return TcpForwardAsync(context, forwarderInfo.EndPoint, linkedCts.Token);
            }

            _logger.LogInfo(context.TraceIdentifier, "Skipping forwarder because protocol is unsupported.");
            return _next(context);
        }

        public async Task UdpForwardAsync(DnsContext context, IPEndPoint server, CancellationToken cancellationToken)
        {
            var request = context.Request;
            var response = context.Response;

            var requestPacket = new DnsPacket(unchecked((ushort)_rand.Next(1, UInt16.MaxValue)), DnsPacketFlag.QueryPacket)
            {
                AcceptNonAuthenticatedData = request.AcceptNonAuthenticatedData,
                RecursionDesired = request.RecursionDesired,
                OpCode = request.OperationCode
            };
            foreach (var item in request.Queries)
            {
                requestPacket.Queries.Add(item);
            }

            DnsPacket responsePacket = null;
            using (var dnsClient = new DnsUdpClient(_packetParser))
            {
                dnsClient.UseRandomTransactionId = true;
                await dnsClient.ConnectAsync(server);

                responsePacket = await dnsClient.SendAsync(requestPacket, cancellationToken);
            }

            if (responsePacket == null)
            {
                _logger.LogInfo(context.TraceIdentifier, "Forwarding failed. Set request to be ignored.");
                response.IgnoreRequest = true;
                return;
            }

            _logger.LogInfo(context.TraceIdentifier, $"Forwarding completed. Reply code is {responsePacket.ReplyCode}.");
            response.ServerAuthoritative = responsePacket.ServerAuthoritative;
            response.AnswerAuthenticated = responsePacket.AnswerAuthenticated;
            response.RecursionAvailable = responsePacket.RecursionAvailable;
            response.ReplyCode = responsePacket.ReplyCode;
            foreach (var item in responsePacket.AnswerRRs)
            {
                response.AnswerRRs.Add(item);
            }
            foreach (var item in responsePacket.AuthoritativeNSs)
            {
                response.AuthoritativeNSs.Add(item);
            }
            foreach (var item in responsePacket.AdditionalRRs)
            {
                response.AdditionalRRs.Add(item);
            }
        }
        public async Task TcpForwardAsync(DnsContext context, IPEndPoint server, CancellationToken cancellationToken)
        {
            var request = context.Request;
            var response = context.Response;

            var requestPacket = new DnsPacket(unchecked((ushort)_rand.Next(1, UInt16.MaxValue)), DnsPacketFlag.QueryPacket)
            {
                AcceptNonAuthenticatedData = request.AcceptNonAuthenticatedData,
                RecursionDesired = request.RecursionDesired,
                OpCode = request.OperationCode
            };
            foreach (var item in request.Queries)
            {
                requestPacket.Queries.Add(item);
            }

            DnsPacket responsePacket = null;
            using (var dnsClient = new DnsTcpClient(_packetParser))
            {
                dnsClient.UseRandomTransactionId = true;
                await dnsClient.ConnectAsync(server);

                responsePacket = await dnsClient.SendAsync(requestPacket, cancellationToken);
            }

            if (responsePacket == null)
            {
                _logger.LogInfo(context.TraceIdentifier, "Forwarding failed. Set request to be ignored.");
                response.IgnoreRequest = true;
                return;
            }

            _logger.LogInfo(context.TraceIdentifier, $"Forwarding completed. Reply code is {responsePacket.ReplyCode}.");
            response.ServerAuthoritative = responsePacket.ServerAuthoritative;
            response.AnswerAuthenticated = responsePacket.AnswerAuthenticated;
            response.RecursionAvailable = responsePacket.RecursionAvailable;
            response.ReplyCode = responsePacket.ReplyCode;
            foreach (var item in responsePacket.AnswerRRs)
            {
                response.AnswerRRs.Add(item);
            }
            foreach (var item in responsePacket.AuthoritativeNSs)
            {
                response.AuthoritativeNSs.Add(item);
            }
            foreach (var item in responsePacket.AdditionalRRs)
            {
                response.AdditionalRRs.Add(item);
            }
        }

        public void Start()
        {
            // nop
        }

        #region IDisposable Support
        private bool disposedValue = false; // 要检测冗余调用

        protected virtual void Dispose(bool disposing)
        {
            if (!disposedValue)
            {
                if (disposing)
                {
                    // 释放托管状态(托管对象)。
                    _disposedToken.Cancel();
                }

                // 释放未托管的资源(未托管的对象)并在以下内容中替代终结器。
                // 将大型字段设置为 null。

                disposedValue = true;
            }
        }

        // 仅当以上 Dispose(bool disposing) 拥有用于释放未托管资源的代码时才替代终结器。
        // ~DnsForwarderMiddleware() {
        //   // 请勿更改此代码。将清理代码放入以上 Dispose(bool disposing) 中。
        //   Dispose(false);
        // }

        // 添加此代码以正确实现可处置模式。
        public void Dispose()
        {
            // 请勿更改此代码。将清理代码放入以上 Dispose(bool disposing) 中。
            Dispose(true);
            // 如果在以上内容中替代了终结器，则取消注释以下行。
            // GC.SuppressFinalize(this);
        }
        #endregion
    }
}
